
#include <nifti/nifti1_io.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <signal.h>
#include <time.h>

#include "graf.h"

float corr(float* data, int offsetx, int offsety, int len) 
{
	float sx = 0.0f;
	float sy = 0.0f;

	for (int i = 0; i < len; i++) 
	{
		sx += data[offsetx + i];
		sy += data[offsety + i];
	}	

	float meanx = sx / len;
	float meany = sy / len;

	float num = 0.0f;
	float denumx = 0.0f;
	float denumy = 0.0f;

	for (int i = 0; i < len; i++) 
	{

		float diffx = data[offsetx + i] - meanx;
		float diffy = data[offsety + i] - meany;

		num += diffx * diffy;

		denumx += diffx * diffx;
		denumy += diffy * diffy;
		
	}

	float denum = sqrt(denumx) * sqrt(denumy);

	if (denum == 0.0f) return 0.0f;
	else return num / denum;

}

struct entry 
{
	unsigned int n1;
	unsigned int n2;
	float r;
};

nifti_image* ni_epi = NULL;
nifti_image* ni_mask = NULL;

float* epi = NULL;
short* mask = NULL;
FILE* fp = NULL;

void quit() 
{
	if (ni_epi != NULL) nifti_image_free(ni_epi);
	if (ni_mask != NULL) nifti_image_free(ni_mask);

    if (epi != NULL) free(epi);
    if (mask != NULL) free(mask);

    if (fp != NULL) fclose(fp);

    exit(0);
}

void quit2(const char *msg)
{
    printf("%s\n", msg);
    quit();
} 

void handler(int signum)
{
    printf("handling signal %d...\n", signum);
    quit();
}

int main(int argc, char** argv) 
{
	if (argc < 4) 
	{
		printf("usage: epi.nii mask.nii threshold adjlist_out\n");
        quit();
	}

    install_handler(handler);

	char* epi_name = argv[1];
	char* mask_name = argv[2];
    float threshold = atof(argv[3]);
	char* fp_name = argv[4];

	ni_epi = nifti_image_read(epi_name, 1);	
    if (ni_epi == NULL) {
        quit2("error reading epi image");
    }

	ni_mask = nifti_image_read(mask_name, 1);	
    if (ni_mask == NULL) {
        quit2("error reading mask image");
    }

    if (ni_epi->datatype != DT_FLOAT) {
        printf("epi datatype is not float");
        quit();
    }

    if (ni_mask->datatype != DT_SIGNED_SHORT) {
        printf("mask datatype is not signed short");
        quit();
    }

	printf("%s and %s loaded.\n", epi_name, mask_name);

	unsigned int ni = ni_epi->dim[1];
	unsigned int nj = ni_epi->dim[2];
	unsigned int nk = ni_epi->dim[3];
	unsigned int nt = ni_epi->dim[4];

	printf("epi is %dx%dx%dx%d\n", ni, nj, nk, nt);

	unsigned int ninj = ni*nj;
	unsigned int ninjnk = ni*nj*nk;
	unsigned int njnk = nj*nk;
	unsigned int njnknt = nj*nk*nt;
	unsigned int nknt = nk*nt;

    // both epi and mask is reorganized, otherwise cpu cache is almost no use!
	printf("starting parallel data reorganization for epi\n");
    
	epi = (float*) calloc(ni*nj*nk*nt, sizeof(float));
    if (epi == NULL) quit2("unable to allocated memory for mask");
	printf("memory for epi allocated.\n");
	float* epi_data = (float*) ni_epi->data;

	#pragma omp parallel
	for (unsigned int i = 0; i < ni; i++) 
		for (unsigned int j = 0; j < nj; j++) 
			for (unsigned int k = 0; k < nk; k++) 
				for (unsigned int t = 0; t < nt; t++) 
					epi[i*njnknt+j*nknt+k*nt+t] = epi_data[i+j*ni+k*ninj+t*ninjnk];

	printf("epi reorganized\n");
	nifti_image_free(ni_epi);
    ni_epi = NULL;
	printf("epi image freed\n");

	printf("starting parallel data reorganization for mask\n");

	mask = (short*) calloc(ni*nj*nk, sizeof(short));
    if (mask == NULL) quit2("unable to allocate memory for mask");
	printf("memory for mask allocated.\n");

	short* mask_data = (short*) ni_mask->data;

	#pragma omp parallel
	for (unsigned int i = 0; i < ni; i++) 
		for (unsigned int j = 0; j < nj; j++) 
			for (unsigned int k = 0; k < nk; k++) 
				mask[i*njnk+j*nk+k] = mask_data[i+j*ni+k*ninj];

	printf("mask reorganized\n");
	nifti_image_free(ni_mask);
    ni_mask = NULL;
	printf("mask image freed\n");

    // we are done with images, just using buffers now
    //
	fp = fopen(fp_name, "wb");
	printf("%s opened\n", fp_name);

	unsigned int num_edges = 0;
	unsigned int num_vertices = 0;

	printf("starting parallel data processing, this will take time....\n");
    unsigned int num_correlations = (ninjnk - 1) * (ninjnk - 2) / 2;
    printf("%d million correlations will be calculated\n", num_correlations/1000000);

	fseek(fp, 24, SEEK_SET);
    
    clock_t start = clock();

    unsigned int cnt = 0;
    unsigned int inside_mask = 0;
	#pragma omp parallel for default(none) shared(ninjnk, nt, num_edges, num_vertices, epi, mask, fp, cnt, threshold, inside_mask, start)
	for (unsigned int offset1 = 0; offset1 < ninjnk; offset1++) 
	{

        #pragma omp critical
        if ((cnt % 1000) == 0)
        {
            printf("%d/%d/%d, V=%d E=%d D=%.4f, elapsed=%ld sec\n", inside_mask, cnt, ninjnk, num_vertices, num_edges, (2.0f*num_edges)/(num_vertices * (num_vertices-1)), (clock() - start)/CLOCKS_PER_SEC);
        }

        #pragma omp atomic
        cnt++;

        if (num_vertices > 1000) continue;

        // omit the ones out of mask
		if (mask[offset1] == 0) continue;

        #pragma omp atomic
        inside_mask++;

        int found = 0;

		for (unsigned int offset2 = offset1+1; offset2 < ninjnk; offset2++) 
		{
            // omit auto correlation
			if (offset1 == offset2) continue;

            // omit the ones out of mask
			if (mask[offset2] == 0) continue;

			float r = corr(epi, offset1*nt, offset2*nt, nt);

			if ((r >= threshold) || (r <= -threshold)) {

                found = 1;

				struct entry e;

				e.n1 = offset1;
				e.n2 = offset2;
				e.r = r;

				#pragma omp critical
				fwrite(&e, sizeof(e), 1, fp);

				#pragma omp atomic
				num_edges++;

			}

		}

        if (found > 0) {
            #pragma omp atomic
            num_vertices++;
        }

	}

	fseek(fp, 0, SEEK_SET);

	fwrite(&ni, sizeof(unsigned int), 1, fp);
	fwrite(&nj, sizeof(unsigned int), 1, fp);
	fwrite(&nk, sizeof(unsigned int), 1, fp);
	fwrite(&nt, sizeof(unsigned int), 1, fp);
	fwrite(&num_vertices, sizeof(unsigned int), 1, fp);
	fwrite(&num_edges, sizeof(unsigned int), 1, fp);

    fclose(fp);
    fp = NULL;

	printf("num_vertices=%d num_edges=%d\n", num_vertices, num_edges);	

    quit();

	return 0;
}
